# before
self.register_buffer(
    "B", torch.empty((self.k // 16, self.n * 16 // 8), dtype=torch.int)
)

# after
self.register_buffer(
    "B", torch.empty((self.k // 16 // 2, self.n * 16 // 8), dtype=torch.int)
)

# before
def mask_creator(tensor):
    ...
    num_groups = tensor.numel() // M

    # N:M sparsity for linear layers
    tensor_temp = tensor.detach().abs().reshape(num_groups, M)

# after
def mask_creator(tensor):
    ...
    num_groups = tensor.numel() // M

    # Subtract the offset value for pruning
    maxq = 2**4 - 1
    ZERO_VALUE = (maxq + 1) // 2
    tensor = tensor - ZERO_VALUE

    # N:M sparsity for linear layers
    tensor_temp = tensor.detach().abs().reshape(num_groups, M)

# before
def pack(self, linear, scales, trans=False):
    ...

    mask = mask_creator(w.T).cuda().bool()
    w = mask * w.T
    w, meta = sparse_semi_structured_from_dense_cutlass(w)
    w = w.t()
# after
def pack(self, linear, scales, trans=False):
    ...

    mask = mask_creator(w.T).cuda().bool()
    # Avoid confusing the pruned elements and the zero elements (-8 + (maxq + 1) // 2)
    w += 1

    w = mask * w.T
    w, meta = sparse_semi_structured_from_dense_cutlass(w)
    
    # Reover the original values
    w -= 1
    w = w.t()